from dotenv import load_dotenv, find_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
import json

_ = load_dotenv(find_dotenv())

@tool
def add(a: int, b: int) -> int:
    """Add two integers.

    Args:
        a: First integer
        b: Second integer
    """
    return a + b

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two integers.

    Args:
        a: First integer
        b: Second integer
    """
    return a * b

llm = ChatOpenAI(model="gpt-4o")
llm_with_tools = llm.bind_tools([add, multiply])

query = "3的4倍是多少?"
messages = [HumanMessage(content=query)]

output = llm_with_tools.invoke(messages)

# print(output)
print(json.dumps(output.tool_calls, indent=4))

messages.append(output)

for tool_call in output.tool_calls:
    selected_tool = {"add":add, "multiply":multiply}[tool_call["name"].lower()]
    tool_msg = selected_tool.invoke(tool_call)
    messages.append(tool_msg)
    
print(messages)
response = llm_with_tools.invoke(messages)
print(response.content)
